{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Cross-attention is extremely simple, same as normal attention but x in Q = W_q(x) and x in W_k(x)/W_v(x) are no longer \n",
    "the same sequence. The \"attention\" paid from queries to keys/vals is now across sequences rather than \n",
    "within a sequence of tokens. This is actually the \"og\" attention in that it's how attention appeared \n",
    "in the original \"Attention is All You Need\" paper which was originally motivated by the language translation task. \n",
    "Cross-attention there allowed any token in \"I am a guy\" to attend to any token in \"Je suis un mec,\" which is obviously \n",
    "something you'd want to do. \n",
    "\n",
    "Cross-attention isn't used this way in LLMs these days (or at all). It is still fundamental because it's used in many other \n",
    "modern settings: RAG to attend to the top-k embeddings a retrieval model/reranker gives you, multimodal LLMs for modality tokens \n",
    "to talk to each other, conditional diffusion models to steer the latent/make the latent a function of an external \n",
    "sequence you want to condition on/use as context, and much more. It's a pretty neat trick to know well if you're planning \n",
    "on reading about/working on arch at all. \n",
    "\n",
    "The interpretation of cross-attention is that the resulting vector (1/D) * SM(Q@K.T) @ V that is \n",
    "a linear combination of values has value vectors now coming from the another sequence, but their weights \n",
    "(co-efficients) are chosen by the SOURCE sequence. \n",
    "\n",
    "Personal note: this is actually an important tool if you're innovating in architectures, for instance \n",
    "one of the first projects I tried years ago attempted to improve the performance of a LM on languaeg (eg. perplexity)\n",
    "by having it attend to embeddings produced by a VISION encoder. The dream was that attending to an image of various \n",
    "flower species would help you distinguish them even when you were presented only textual descriptions. \n",
    "Ultimately this didn't work but others got a related idea using cross-attention to work augmenting \n",
    "LLMs with skills that other LLMs have, see (https://arxiv.org/pdf/2401.02412).\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn \n",
    "import torch.nn.functional as F \n",
    "import math \n",
    "\n",
    "# let's quickly reimplement mhsa as a baseline that we can then tweak. \n",
    "# good practice to implement this from scratch many times!\n",
    "class MHSA(nn.Module): \n",
    "    def __init__(self, head_dim=64, D=512, causal=True): \n",
    "        super().__init__()\n",
    "        self.head_dim=head_dim\n",
    "        self.D = D\n",
    "        assert D % head_dim == 0, \"Error: head dimension does not evenly divide residual stream dimension.\"\n",
    "        self.nheads = D//head_dim\n",
    "        self.causal = causal \n",
    "        \n",
    "        self.qkv_proj = nn.Linear(D, 3 * D)\n",
    "        self.out_proj = nn.Linear(D, D)\n",
    "\n",
    "    def forward(self, x): # bsd -> bsd \n",
    "        B, S, _ = x.shape\n",
    "        qkv = self.qkv_proj(x)  # [b, s, 3*D]\n",
    "        qkv = qkv.reshape(B, S, 3, self.nheads, self.head_dim)\n",
    "        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, b, nheads, s, head_dim]\n",
    "        q, k, v = qkv[0], qkv[1], qkv[2]  # each [b, nheads, s, head_dim]\n",
    "\n",
    "        attn_logits = torch.einsum('bnid,bnjd->bnij', q, k) # output is [b, nh, s, s]\n",
    "        normalize = math.sqrt(self.head_dim)\n",
    "        attn_logits = attn_logits / normalize # elementwise \n",
    "\n",
    "        if self.causal: \n",
    "            mask = torch.arange(S)[:, None] >= torch.arange(S) # [s, s]\n",
    "            attn_logits = torch.where(mask, attn_logits, float('-inf') * torch.ones_like(attn_logits))\n",
    "\n",
    "        A = F.softmax(attn_logits, dim=-1) # [b, nh, s, s]\n",
    "\n",
    "        # out = A @ V then cat all attn heads together to get a long value vector of dim D\n",
    "        out = torch.einsum('bnij,bnjd->bnid', A, v) # [b, nh, s, s] @ [b, nh, s, hd] -> [b, nh, s, hd]\n",
    "        out = out.transpose(1, 2).reshape(B, S, -1) # [b, s, d]\n",
    "\n",
    "        return self.out_proj(out) # [d, d] @ [b, s, d] along last dim --> [b, s, d] outputs \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CrossMHSA(nn.Module): \n",
    "    def __init__(self, head_dim=64, D=512): \n",
    "        super().__init__()\n",
    "        self.head_dim = head_dim\n",
    "        self.D = D\n",
    "        assert D % head_dim == 0, \"Error: head dimension does not evenly divide residual stream dimension.\"\n",
    "        self.nheads = D//head_dim\n",
    "        \n",
    "        self.q_proj = nn.Linear(D, D)\n",
    "        self.k_proj = nn.Linear(D, D)\n",
    "        self.v_proj = nn.Linear(D, D)\n",
    "        self.out_proj = nn.Linear(D, D)\n",
    "\n",
    "    def forward(self, x1, x2): # (bs1d, bs2d) -> bsd\n",
    "        B, S1, _ = x1.shape\n",
    "        _, S2, _ = x2.shape\n",
    "        q = self.q_proj(x1) # [b, s1, d]\n",
    "        k, v = self.k_proj(x2), self.v_proj(x2)\n",
    "\n",
    "        q = q.reshape(B, S1, self.nheads, self.head_dim).transpose(1, 2)  # [b, nh, s1, hd]\n",
    "        k = k.reshape(B, S2, self.nheads, self.head_dim).transpose(1, 2)  # [b, nh, s2, hd]\n",
    "        v = v.reshape(B, S2, self.nheads, self.head_dim).transpose(1, 2)  # [b, nh, s2, hd]\n",
    "\n",
    "        attn_logits = torch.einsum('bnid,bnjd->bnij', q, k) # output is [b, nh, s1, s2]\n",
    "        normalize = math.sqrt(self.head_dim)\n",
    "        attn_logits = attn_logits / normalize # elementwise \n",
    "\n",
    "        A = F.softmax(attn_logits, dim=-1) # [b, nh, s1, s2]\n",
    "\n",
    "        # out = A @ V then cat all attn heads together to get a long value vector of dim D\n",
    "        out = torch.einsum('bnij,bnjd->bnid', A, v) # [b, nh, s1, s2] @ [b, nh, s2, hd] -> [b, nh, s1, hd]\n",
    "        out = out.transpose(1, 2).reshape(B, S1, -1) # [b, s1, d]\n",
    "\n",
    "        return self.out_proj(out) # [d, d] @ [b, s, d] along last dim --> [b, s, d] outputs \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when sequences are identical: outputs match = True\n",
      "cross-attention allows different sequence lengths, while self-attention requires identical sequences\n",
      "cross-attention gives different results with different sequences: True\n",
      "all tests completed!\n"
     ]
    }
   ],
   "source": [
    "# tests\n",
    "batch_size = 2\n",
    "seq_len = 5\n",
    "dim = 512\n",
    "head_dim = 64\n",
    "\n",
    "# test 1: when sequences are the same\n",
    "x = torch.randn(batch_size, seq_len, dim)\n",
    "\n",
    "# initialize both attention modules\n",
    "# for fair comparison, we need to set causal=False in MHSA\n",
    "mhsa = MHSA(head_dim=head_dim, D=dim, causal=False)\n",
    "cross_attn = CrossMHSA(head_dim=head_dim, D=dim)\n",
    "\n",
    "# initialize weights to be the same\n",
    "with torch.no_grad():\n",
    "    # set qkv projections to be equivalent\n",
    "    mhsa.qkv_proj.weight.copy_(torch.cat([\n",
    "        cross_attn.q_proj.weight,\n",
    "        cross_attn.k_proj.weight,\n",
    "        cross_attn.v_proj.weight\n",
    "    ], dim=0))\n",
    "    mhsa.qkv_proj.bias.copy_(torch.cat([\n",
    "        cross_attn.q_proj.bias,\n",
    "        cross_attn.k_proj.bias,\n",
    "        cross_attn.v_proj.bias\n",
    "    ], dim=0))\n",
    "    \n",
    "    # set output projections to be the same\n",
    "    cross_attn.out_proj.weight.copy_(mhsa.out_proj.weight)\n",
    "    cross_attn.out_proj.bias.copy_(mhsa.out_proj.bias)\n",
    "\n",
    "# forward pass\n",
    "mhsa_output = mhsa(x)\n",
    "cross_attn_output = cross_attn(x, x)  # same sequence for both inputs\n",
    "\n",
    "# check if outputs are the same when inputs are the same\n",
    "is_close = torch.allclose(mhsa_output, cross_attn_output, rtol=1e-4, atol=1e-4)\n",
    "print(f\"when sequences are identical: outputs match = {is_close}\")\n",
    "\n",
    "# test 2: when sequences are different\n",
    "seq_len1 = 5\n",
    "seq_len2 = 7\n",
    "x1 = torch.randn(batch_size, seq_len1, dim)\n",
    "x2 = torch.randn(batch_size, seq_len2, dim)\n",
    "\n",
    "# forward pass with different sequences\n",
    "cross_attn_diff = cross_attn(x1, x2)\n",
    "\n",
    "# check output shape\n",
    "expected_shape = (batch_size, seq_len1, dim)\n",
    "assert cross_attn_diff.shape == expected_shape, f\"expected shape {expected_shape}, got {cross_attn_diff.shape}\"\n",
    "\n",
    "# try to compute regular self-attention with different sequence lengths (this should fail or give different results)\n",
    "try:\n",
    "    # this won't work directly with MHSA as it expects a single sequence\n",
    "    print(\"cross-attention allows different sequence lengths, while self-attention requires identical sequences\")\n",
    "except Exception as e:\n",
    "    print(f\"as expected, self-attention can't handle different sequence lengths: {e}\")\n",
    "\n",
    "# verify that cross-attention with different sequences gives different results than with identical sequences\n",
    "x1_copy = x1.clone()\n",
    "cross_attn_same = cross_attn(x1_copy, x1_copy)\n",
    "is_different = not torch.allclose(cross_attn_same, cross_attn_diff, rtol=1e-4, atol=1e-4)\n",
    "print(f\"cross-attention gives different results with different sequences: {is_different}\")\n",
    "\n",
    "print(\"all tests completed!\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "envi",
   "language": "python",
   "name": "lingua_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
